import os,sys

tasks = os.listdir("D:/ai/t5/t5_chat/save")

loras = []
for task in tasks:
    for lora in os.listdir("D:/ai/t5/t5_chat/save/"+task):
        loras.append(task+"/"+lora)
print(loras)
#%%
print(len(loras))

ner_li = []
re_li = []
ee_li = []

for lora in loras:
    if "ner" in lora :
        ner_li.append(lora)
    elif "re" in lora:
        re_li.append(lora)
    elif "ee" in lora:
        ee_li.append(lora)
print(ner_li)
print(re_li)
print(ee_li)

#%%
import re
from typing import List, Tuple
def parse_tuple_string(tuple_string: str) -> List[Tuple[str, str]]:
    # 使用正则表达式匹配元组，使第二个元素可选
    pattern = re.compile(r'\(([^,]+),? ?([^)]+)?\)')
    return pattern.findall(tuple_string)

def test_parse_tuple_string():
    tuple_string = "[(bankruptcy, declare bankruptcy)]"
    print(parse_tuple_string(tuple_string))
    print(parse_tuple_string("[()]"))
    print(parse_tuple_string("[(bankruptcy, declare bankruptcy)"))
    print(parse_tuple_string("[(钢铁, 单位)]"))
    print(parse_tuple_string("[(钢铁,)]"))
    print(parse_tuple_string("[(, 单位), (铁路, 单位)]"))
#%%
import re
from typing import List, Tuple

def re_calculate_metrics(predicted: List[str], actual: List[str]):
    # 将字符串转换为元组
    def re_parse_tuple_string(tuple_string: str) -> List[Tuple[str, str, str]]:
        # 使用正则表达式匹配元组，使第二个和第三个元素可选
        tuple_string = tuple_string.replace('，',',').replace('（','(').replace('）', ')').replace("：", ":").replace(" ", "")
        pattern = re.compile(r'\(([^,]*),? ?([^,]*)?,? ?([^)]*)?\)')
        tuples = pattern.findall(tuple_string)
        # 将空元组替换为特殊元组
        tuples = [('EMPTY', 'EMPTY', 'EMPTY') if t == ('', '', '') else t for t in tuples]
        return tuples

    predicted = [re_parse_tuple_string(t) for t in predicted]
    actual = [re_parse_tuple_string(t) for t in actual]

    # 将每个列表中的元组抽取出来
    predicted_tuples = [item for sublist in predicted for item in sublist]
    actual_tuples = [item for sublist in actual for item in sublist]

    tp = len(set(predicted_tuples) & set(actual_tuples))
    fp = len(set(predicted_tuples) - set(actual_tuples))
    fn = len(set(actual_tuples) - set(predicted_tuples))

    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

    return precision, recall, f1

def ner_calculate_metrics(predicted: List[str], actual: List[str]):
    # 将字符串转换为元组
    def ner_parse_tuple_string(tuple_string: str) -> List[Tuple[str, str]]:
        # 使用正则表达式匹配元组，使第二个元素可选
        tuple_string = tuple_string.replace('，',',').replace('（','(').replace('）', ')').replace("：", ":").replace(" ", "")
        pattern = re.compile(r'\(([^,]*),? ?([^)]*)?\)')
        tuples = pattern.findall(tuple_string)
        # 将空元组替换为特殊元组
        tuples = [('EMPTY', 'EMPTY') if t == ('', '') else t for t in tuples]
        return tuples

    predicted = [ner_parse_tuple_string(t) for t in predicted]
    actual = [ner_parse_tuple_string(t) for t in actual]

    # 将每个列表中的元组抽取出来
    predicted_tuples = [item for sublist in predicted for item in sublist]
    actual_tuples = [item for sublist in actual for item in sublist]

    tp = len(set(predicted_tuples) & set(actual_tuples))
    fp = len(set(predicted_tuples) - set(actual_tuples))
    fn = len(set(actual_tuples) - set(predicted_tuples))

    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

    return precision, recall, f1

def ee_calculate_metrics(predicted: List[str], actual: List[str]):
    def ee_parse_tuple_string(tuple_string: str) -> List[Tuple[str, str, List[Tuple[str, str]]]]:
        def parse_tuple_string(tuple_string: str):
            # 使用正则表达式匹配元组，使第二个元素可选
            tuple_string = tuple_string.replace('，',',').replace('（','(').replace('）', ')').replace("：", ":").replace(" ", "")
            pattern = re.compile(r'\(([^,]*),? ?([^)]*)?\)')
            tuples = pattern.findall(tuple_string)
            # 将空元组替换为特殊元组
            tuples = tuple([('EMPTY', 'EMPTY') if t == ('', '') else t for t in tuples])
            return tuples
        # 使用正则表达式匹配元组，使第二个和第三个元素可选
        tuple_string = tuple_string.replace('，',',').replace('（','(').replace('）', ')').replace("：", ":").replace(" ", "")
        pattern = re.compile(r'\(([^,]*),? ?([^,]*)?,? ?(\[.*\])?\)')
        tuples = pattern.findall(tuple_string)
        # 将空元组替换为特殊元组
        tuples = [('EMPTY', 'EMPTY', [('EMPTY', 'EMPTY')]) if t == ('', '', '') else (t[0], t[1], tuple(sorted(parse_tuple_string(t[2])))) for t in tuples]
        return tuples
    # 将字符串转换为元组
    predicted = [ee_parse_tuple_string(t) for t in predicted]
    actual = [ee_parse_tuple_string(t) for t in actual]

    # 将每个列表中的元组抽取出来
    predicted_tuples = [item for sublist in predicted for item in sublist]
    actual_tuples = [item for sublist in actual for item in sublist]

    tp = len(set(predicted_tuples) & set(actual_tuples))
    fp = len(set(predicted_tuples) - set(actual_tuples))
    fn = len(set(actual_tuples) - set(predicted_tuples))

    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

    return precision, recall, f1

def test_calculate_metrics():
    # print(tuple([('1', '1'), ('2', '2')]))
    # print(ner_calculate_metrics(["[(bankruptcy, declare bankruptcy)]"], ["[(bankruptcy, declare bankruptcy)]"]))
    # print(ner_calculate_metrics(["[(bankruptcy, declare bankruptcy)]"], ["[(bankruptcy, declare bankruptcy), (bankruptcy, declare bankruptcy)]"]))
    # print(ner_calculate_metrics(["[()]"], ["[()]"]))
    print(ee_calculate_metrics(["[(bankruptcy, declare bankruptcy, [(bankruptcy, declare bankruptcy)])]"], ["[(bankruptcy, declare bankruptcy, [(bankruptcy, declare bankruptcy)])]"]))
    print(ee_calculate_metrics(["[(bankruptcy, declare bankruptcy, [(1, 1), (2, 2)])]"], ["[(bankruptcy, declare bankruptcy, [(1, 1), (2, 2)])]"]))